from flask import Flask, request
from core.db.dbeaver_pool import DbeaverPool
import json


class ProxyApi(object):
    def __init__(self):
        self.app = Flask(__name__)
        self.dbeaver_pool = DbeaverPool()

        @self.app.route('/random')
        def random():
            # 获取参数
            protocol = request.args.get('protocol')
            domain = request.args.get('domain')
            proxy = self.dbeaver_pool.random_proxy(protocol=protocol, domain=domain)
            if protocol:
                return f'{protocol}://{proxy.ip}:{proxy.port}'
            else:
                return f'{proxy.ip}:{proxy.port}'

        @self.app.route('/proxies')
        def proxies():
            protocol = request.args.get('protocol')
            domain = request.args.get('domain')
            proxy_list = self.dbeaver_pool.get_proxies(protocol=protocol, domain=domain)
            proxy_list = [proxy.__dict__ for proxy in proxy_list]
            return json.dumps(proxy_list, ensure_ascii=False)

        @self.app.route('/disable_domain')
        def disable_domain():
            ip = request.args.get('ip')
            domain = request.args.get('domain')
            if ip is None:
                return '请提供ip参数'
            if domain is None:
                return '请提供域名domain参数'
            self.dbeaver_pool.add_disable_domain(ip, domain)
            return f'ip: {ip} 禁用域名 {domain} 成功'

    def run(self):
        self.app.run('0.0.0.0', port=16888)

    @classmethod
    def start(cls):
        proxy_api = cls()
        proxy_api.run()


if __name__ == '__main__':
    ProxyApi.start()
